import torch
import os
import torch.distributed as dist
from torch.utils.data.distributed import DistributedSampler
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    BitsAndBytesConfig,
    TrainingArguments
)
from peft import LoraConfig, prepare_model_for_kbit_training
from trl import SFTTrainer
from data.dataset_utils import create_code_review_dataset

# 模型配置
model_name = "meta-llama/Llama-2-7b-hf"
model_name = "/root/autodl-tmp/Llama-2-7b-hf/"
tokenizer = AutoTokenizer.from_pretrained(model_name,trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True
)

# 初始化分布式环境
if dist.is_available():
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ.get("LOCAL_RANK", -1))
    torch.cuda.set_device(local_rank)

# 修改模型加载
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    trust_remote_code=True,
    torch_dtype=torch.float16
).to(f"cuda:{local_rank}" if local_rank != -1 else "cuda")
model = prepare_model_for_kbit_training(model)

# LoRA配置
peft_config = LoraConfig(
    r=8,
    lora_alpha=32,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)

# 数据加载
dataset = create_code_review_dataset()
train_data = dataset["train"]
eval_data = dataset["test"]

# 训练参数
training_args = TrainingArguments(
    output_dir="./qlora-output",
    num_train_epochs=3,
    per_device_train_batch_size=2,
    gradient_accumulation_steps=4,
    optim="paged_adamw_8bit",
    save_steps=500,
    logging_steps=50,
    learning_rate=2e-5,
    fp16=True,
    evaluation_strategy="steps",
    eval_steps=100,
    max_grad_norm=1.0,
    deepspeed="./configs/ds_config.json",
    report_to="wandb"
)

def formatting_prompts_func(examples):
    output_texts = []
    for i in range(len(examples['instruction'])):
        text = f"instruction：{examples['instruction'][i]}\ninput：{examples['input'][i]}\noutput：{examples['output'][i]}"
        output_texts.append(text)
    return output_texts

# 训练器配置
trainer = SFTTrainer(
    model=model,
    train_dataset=train_data,
    eval_dataset=eval_data,
    peft_config=peft_config,
    formatting_func=formatting_prompts_func,
    tokenizer=tokenizer,
    args=training_args,
)

# 开始训练
trainer.train()
